#!/bin/bash

# 优化版URPC2020训练脚本
# 基于训练日志分析的改进建议

echo "======================================================================"
echo "🌊 URPC2020水下目标检测 - 优化训练"
echo "======================================================================"
echo ""

# 显示优化说明
echo "🎯 本次优化针对以下问题:"
echo "   ❌ 整体Recall偏低 (64.9%)"
echo "   ❌ Scallop类别Recall极低 (42.7%)"
echo "   ❌ mAP@0.5-0.95偏低 (46.5%) - 边界框定位不精确"
echo ""
echo "✅ 主要改进措施:"
echo "   ✓ 保持YOLOv8n模型 (轻量级，快速训练)"
echo "   ✓ 图像尺寸: 640 -> 800px"
echo "   ✓ Early Stopping patience: 25 -> 50"
echo "   ✓ 保持训练轮数: 120 epochs"
echo "   ✓ 增强数据增强策略"
echo "   ✓ 提高Box loss权重"
echo "   ✓ 降低验证置信度阈值"
echo "   ✓ 类别权重优化"
echo ""

# 检查数据集
echo "📊 检查URPC2020数据集..."
if [ ! -d "datasets/urpc2020" ]; then
    echo "❌ 数据集目录不存在: datasets/urpc2020"
    exit 1
fi

if [ ! -f "urpc2020.yaml" ]; then
    echo "❌ 数据集配置文件不存在: urpc2020.yaml"
    exit 1
fi

# 统计数据集
TRAIN_COUNT=$(ls -1 datasets/urpc2020/train/images/*.jpg 2>/dev/null | wc -l)
VALID_COUNT=$(ls -1 datasets/urpc2020/valid/images/*.jpg 2>/dev/null | wc -l)
TEST_COUNT=$(ls -1 datasets/urpc2020/test/images/*.jpg 2>/dev/null | wc -l)

echo "✅ 数据集已就绪:"
echo "   训练集: $TRAIN_COUNT 张图像"
echo "   验证集: $VALID_COUNT 张图像"
echo "   测试集: $TEST_COUNT 张图像"
echo ""

# 检查GPU
echo "🖥️  检查GPU状态..."
nvidia-smi --query-gpu=index,name,memory.total,memory.free --format=csv,noheader,nounits | while read line; do
    echo "   GPU: $line"
done
echo ""

# 检查Python环境
echo "🐍 Python环境:"
python --version
echo ""

# 检查YOLO库
echo "📦 检查Ultralytics YOLO..."
python -c "from ultralytics import YOLO; import ultralytics; print(f'✅ Ultralytics {ultralytics.__version__} 已安装')" 2>&1
if [ $? -ne 0 ]; then
    echo "❌ Ultralytics YOLO 未安装"
    echo "💡 安装命令: pip install ultralytics"
    exit 1
fi
echo ""

# 训练参数
MODEL="yolov8n.pt"          # 使用YOLOv8-Nano (轻量级)
DATA="urpc2020.yaml"
EPOCHS=120                   # 训练轮数
BATCH=16                     # 可根据显存调整
IMGSZ=800                    # 增加图像尺寸
DEVICE="0"
USE_CLASS_WEIGHTS="--use-class-weights"  # 启用类别权重

echo "⚙️  优化训练参数:"
echo "   模型: $MODEL (保持轻量级)"
echo "   数据: $DATA"
echo "   轮数: $EPOCHS epochs"
echo "   批次: $BATCH"
echo "   尺寸: ${IMGSZ}px (↑ from 640)"
echo "   设备: GPU $DEVICE"
echo "   类别权重: 启用 (Scallop加权)"
echo ""

# 预估训练时间
echo "⏱️  预估训练时间:"
echo "   基于RTX 4090, 约2.0-2.5小时"
echo "   (取决于实际收敛速度)"
echo ""

# 确认开始
echo "======================================================================"
echo "🚀 准备开始优化训练..."
echo "======================================================================"
echo ""
echo "💡 训练过程中可以:"
echo "   - 查看实时日志"
echo "   - 使用Ctrl+C安全停止"
echo "   - 使用tensorboard监控: tensorboard --logdir runs/train"
echo ""

read -p "按Enter键开始训练，或Ctrl+C取消..." -t 10 || echo ""

echo ""
echo "======================================================================"
echo "🏋️ 开始训练..."
echo "======================================================================"
echo ""

# 记录开始时间
START_TIME=$(date +%s)

# 运行训练
python train_urpc2020_optimized.py \
    --model $MODEL \
    --data $DATA \
    --epochs $EPOCHS \
    --batch $BATCH \
    --imgsz $IMGSZ \
    --device $DEVICE \
    --project runs/train \
    --name urpc2020_optimized \
    $USE_CLASS_WEIGHTS \
    --exist-ok

# 记录结束时间
END_TIME=$(date +%s)
DURATION=$((END_TIME - START_TIME))
HOURS=$((DURATION / 3600))
MINUTES=$(((DURATION % 3600) / 60))

echo ""
echo "======================================================================"
echo "✅ 训练脚本执行完成"
echo "======================================================================"
echo ""
echo "⏱️  训练用时: ${HOURS}小时 ${MINUTES}分钟"
echo ""
echo "📝 查看结果:"
echo "   训练目录: runs/train/urpc2020_optimized/"
echo "   最佳模型: runs/train/urpc2020_optimized/weights/best.pt"
echo "   最终模型: runs/train/urpc2020_optimized/weights/last.pt"
echo "   训练曲线: runs/train/urpc2020_optimized/results.png"
echo "   混淆矩阵: runs/train/urpc2020_optimized/confusion_matrix.png"
echo ""
echo "📊 性能对比:"
echo "   运行对比脚本: python compare_models.py"
echo ""
echo "🔬 进一步验证:"
echo "   1. TensorBoard: tensorboard --logdir runs/train/urpc2020_optimized"
echo "   2. 测试集验证: yolo val model=runs/train/urpc2020_optimized/weights/best.pt data=urpc2020.yaml"
echo "   3. 可视化预测: yolo predict model=runs/train/urpc2020_optimized/weights/best.pt source=datasets/urpc2020/test/images conf=0.20"
echo ""
echo "💡 如果性能提升明显，可以继续优化:"
echo "   - 尝试更大模型: yolov8m.pt 或 yolov8l.pt"
echo "   - 增加图像尺寸: 960 或 1024"
echo "   - 使用模型集成 (ensemble)"
echo ""

